{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RigidBodySim.jl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RigidBodySim is a simulation environment built on top of [RigidBodyDynamics.jl](https://github.com/JuliaRobotics/RigidBodyDynamics.jl), [DifferentialEquations.jl](https://github.com/JuliaDiffEq/DifferentialEquations.jl), and [MeshCat.jl](https://github.com/rdeits/MeshCat.jl).\n",
    "\n",
    "See the [latest documentation](https://JuliaRobotics.github.io/RigidBodySim.jl/latest) and the [quick start guide](https://nbviewer.jupyter.org/github/JuliaRobotics/RigidBodySim.jl/blob/master/notebooks/Quick%20start%20guide.ipynb) for more information and examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading the Atlas `Mechanism`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll first create a `RigidBodyDynamics.Mechanism` for the Atlas robot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WebIO: setting up mux\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "    <script class='js-collapse-script'>\n",
       "        var curMatch =\n",
       "            window.location.href\n",
       "            .match(/(.*?)\\/notebooks\\/.*\\.ipynb/);\n",
       "\n",
       "        curMatch = curMatch ||\n",
       "            window.location.href\n",
       "            .match(/(.*?)\\/apps\\/.*\\.ipynb/);\n",
       "\n",
       "        if ( curMatch ) {\n",
       "            $('head').append('<base href=\"' + curMatch[1] + '/\">');\n",
       "        }\n",
       "    </script>\n"
      ],
      "text/plain": [
       "HTML{String}(\"    <script class='js-collapse-script'>\\n        var curMatch =\\n            window.location.href\\n            .match(/(.*?)\\\\/notebooks\\\\/.*\\\\.ipynb/);\\n\\n        curMatch = curMatch ||\\n            window.location.href\\n            .match(/(.*?)\\\\/apps\\\\/.*\\\\.ipynb/);\\n\\n        if ( curMatch ) {\\n            \\$('head').append('<base href=\\\"' + curMatch[1] + '/\\\">');\\n        }\\n    </script>\\n\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script class='js-collapse-script' src='pkg/WebIO/webio/dist/bundle.js'></script>"
      ],
      "text/plain": [
       "HTML{String}(\"<script class='js-collapse-script' src='pkg/WebIO/webio/dist/bundle.js'></script>\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script class='js-collapse-script' src='pkg/WebIO/providers/ijulia_setup.js'></script>"
      ],
      "text/plain": [
       "HTML{String}(\"<script class='js-collapse-script' src='pkg/WebIO/providers/ijulia_setup.js'></script>\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "  <script class='js-collapse-script'>\n",
       "    $('.js-collapse-script').parent('.output_subarea').css('padding', '0');\n",
       "  </script>\n"
      ],
      "text/plain": [
       "HTML{String}(\"  <script class='js-collapse-script'>\\n    \\$('.js-collapse-script').parent('.output_subarea').css('padding', '0');\\n  </script>\\n\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WebIO: setting up ijulia\n",
      "WebIO: setting up atom\n"
     ]
    }
   ],
   "source": [
    "using RigidBodySim\n",
    "using RigidBodyDynamics\n",
    "using AtlasRobot # tiny package containing Atlas URDF and meshes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Spanning tree:\n",
       "Vertex: world (root)\n",
       "  Vertex: ltorso, Edge: back_bkz\n",
       "    Vertex: mtorso, Edge: back_bky\n",
       "      Vertex: utorso, Edge: back_bkx\n",
       "        Vertex: l_clav, Edge: l_arm_shz\n",
       "          Vertex: l_scap, Edge: l_arm_shx\n",
       "            Vertex: l_uarm, Edge: l_arm_ely\n",
       "              Vertex: l_larm, Edge: l_arm_elx\n",
       "                Vertex: l_ufarm, Edge: l_arm_uwy\n",
       "                  Vertex: l_lfarm, Edge: l_arm_mwx\n",
       "                    Vertex: l_hand, Edge: l_arm_lwy\n",
       "        Vertex: head, Edge: neck_ay\n",
       "        Vertex: r_clav, Edge: r_arm_shz\n",
       "          Vertex: r_scap, Edge: r_arm_shx\n",
       "            Vertex: r_uarm, Edge: r_arm_ely\n",
       "              Vertex: r_larm, Edge: r_arm_elx\n",
       "                Vertex: r_ufarm, Edge: r_arm_uwy\n",
       "                  Vertex: r_lfarm, Edge: r_arm_mwx\n",
       "                    Vertex: r_hand, Edge: r_arm_lwy\n",
       "  Vertex: l_uglut, Edge: l_leg_hpz\n",
       "    Vertex: l_lglut, Edge: l_leg_hpx\n",
       "      Vertex: l_uleg, Edge: l_leg_hpy\n",
       "        Vertex: l_lleg, Edge: l_leg_kny\n",
       "          Vertex: l_talus, Edge: l_leg_aky\n",
       "            Vertex: l_foot, Edge: l_leg_akx\n",
       "  Vertex: r_uglut, Edge: r_leg_hpz\n",
       "    Vertex: r_lglut, Edge: r_leg_hpx\n",
       "      Vertex: r_uleg, Edge: r_leg_hpy\n",
       "        Vertex: r_lleg, Edge: r_leg_kny\n",
       "          Vertex: r_talus, Edge: r_leg_aky\n",
       "            Vertex: r_foot, Edge: r_leg_akx\n",
       "No non-tree joints."
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "atlas = AtlasRobot.mechanism(floating=false)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MechanismState{Float64, Float64, Float64, …}(…)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "const state = MechanismState(atlas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting up visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll again use `Meshcat` for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "using MeshCatMechanisms\n",
    "using RigidBodySim.Visualization.MeshCatInterface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "visuals = URDFVisuals(AtlasRobot.urdfpath();\n",
    "    package_path = [AtlasRobot.packagepath()])\n",
    "visualizer = MechanismVisualizer(atlas, visuals);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Listening on 0.0.0.0:8700...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[36mINFO: \u001b[39m\u001b[22m\u001b[36mServing MeshCat visualizer at http://127.0.0.1:8700\n",
      "\u001b[39m"
     ]
    }
   ],
   "source": [
    "window(visualizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closed loop dynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's simulate Atlas with a very simple controller:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "control! (generic function with 1 method)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function control!(τ, t, state)\n",
    "    # just add a bit of joint damping\n",
    "    τ .= -0.1 .* velocity(state)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "closed_loop_dynamics = Dynamics(atlas, control!);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or you can set up the controller to run at a fixed rate (in terms of simulation time):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# τ = similar(velocity(state))\n",
    "# Δt = 1e-3\n",
    "# controller = PeriodicController(similar(velocity(state)), Δt, control!)\n",
    "# closed_loop_dynamics = Dynamics(atlas, controller);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating an `ODEProblem`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An `ODEProblem` (from the DifferentialEquations.jl ecosystem) represents a to-be-integrated ordinary differential equation (ODE):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DiffEqBase.ODEProblem with uType Array{Float64,1} and tType Float64. In-place: true\n",
       "timespan: (0.0, 5.0)\n",
       "u0: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zero!(state)\n",
    "timespan = (0., 5.)\n",
    "problem = ODEProblem(closed_loop_dynamics, state, timespan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Integrating the ODE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To integrate the ODE, we'll be using the Tsitouras 5/4 (5th order) variable time step Runge-Kutta method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrdinaryDiffEq.Tsit5()"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "integrator = Tsit5()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To visualize Atlas while simulating, we'll create a callback that will be run periodically during the integration process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualizer_callback = CallbackSet(visualizer, state);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can solve the `ODEProblem`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "sol = solve(problem,\n",
    "    integrator,\n",
    "    abs_tol = 1e-10,\n",
    "    dt = 0.05,\n",
    "    callback = CallbackSet(visualizer_callback));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the solution after the fact:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "RigidBodySim.animate(visualizer, state, sol)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.6.2",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
